#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
land30m_downloader.py

Land30m 全球土地利用数据下载模块

功能：
- 从 CAS Earth 数据平台下载 Land30m TIF 文件
- 支持断点续传和进度显示
- 自动创建下载目录
- 错误处理和重试机制
- 配置文件支持

依赖：requests, tqdm

用法示例：
1) 基本下载：
   python land30m_downloader.py --username your_username --output-dir ./data

2) 使用配置文件：
   python land30m_downloader.py --config config.json

3) 仅列出文件不下载：
   python land30m_downloader.py --username your_username --list-only
"""

import argparse
import json
import os
import time
from pathlib import Path
from typing import Dict, List, Optional
from urllib.parse import unquote

import requests
from tqdm import tqdm


class Land30mDownloader:
    """Land30m 数据下载器"""
    
    def __init__(self, username: str, output_dir: str, max_retries: int = 3, timeout: int = 30):
        """
        初始化下载器
        
        Args:
            username: CAS Earth 用户名
            output_dir: 下载目录
            max_retries: 最大重试次数
            timeout: 请求超时时间（秒）
        """
        self.username = username
        self.output_dir = Path(output_dir)
        self.max_retries = max_retries
        self.timeout = timeout
        
        # API 端点
        self.api_url = "http://data.casearth.cn/api/getAllFileListBySdoId?sdoId=6123651428a58f70c2a51e49"
        self.download_base_url = "http://data.casearth.cn/sdo/downloadOneFile"
        
        # 创建下载目录
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # 会话对象，支持连接复用
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
        })
    
    def get_file_list(self) -> List[Dict]:
        """
        获取文件列表
        
        Returns:
            文件信息列表
        """
        print("正在获取文件列表...")
        try:
            response = self.session.get(self.api_url, timeout=self.timeout)
            response.raise_for_status()
            
            data = response.json()
            file_list = data.get("文件信息列表", [])
            
            print(f"找到 {len(file_list)} 个文件")
            return file_list
            
        except requests.exceptions.RequestException as e:
            print(f"获取文件列表失败: {e}")
            return []
        except json.JSONDecodeError as e:
            print(f"解析文件列表失败: {e}")
            return []
    
    def sanitize_filename(self, filename: str) -> str:
        """
        清理文件名，移除非法字符
        
        Args:
            filename: 原始文件名
            
        Returns:
            清理后的文件名
        """
        # Windows 文件名非法字符
        illegal_chars = '<>:"/\\|?*'
        for char in illegal_chars:
            filename = filename.replace(char, '_')
        
        # 移除前后空格和点
        filename = filename.strip(' .')
        
        # 确保文件名不为空
        if not filename:
            filename = "unknown_file"
        
        return filename
    
    def download_file(self, file_info: Dict, progress_bar: Optional[tqdm] = None) -> bool:
        """
        下载单个文件
        
        Args:
            file_info: 文件信息字典
            progress_bar: 进度条对象
            
        Returns:
            下载是否成功
        """
        file_id = file_info.get("id")
        filename = file_info.get("filename", "unknown")
        
        if not file_id:
            print(f"文件 {filename} 缺少 ID，跳过")
            return False
        
        # 清理文件名
        safe_filename = self.sanitize_filename(filename)
        file_path = self.output_dir / safe_filename
        
        # 检查文件是否已存在
        if file_path.exists():
            print(f"文件 {safe_filename} 已存在，跳过")
            if progress_bar:
                progress_bar.update(1)
            return True
        
        # 构建下载 URL
        download_url = f"{self.download_base_url}?id={file_id}&username={self.username}"
        print(f"下载 URL: {download_url}")
        
        # 重试下载
        for attempt in range(self.max_retries):
            try:
                print(f"正在下载 {safe_filename} (尝试 {attempt + 1}/{self.max_retries})")
                
                response = self.session.get(download_url, timeout=self.timeout, stream=True)
                response.raise_for_status()
                
                # 获取文件大小
                total_size = int(response.headers.get('content-length', 0))
                
                # 写入文件
                with open(file_path, 'wb') as f:
                    if total_size > 0:
                        # 使用进度条
                        with tqdm(total=total_size, unit='B', unit_scale=True, 
                                desc=safe_filename, leave=False) as pbar:
                            for chunk in response.iter_content(chunk_size=8192):
                                if chunk:
                                    f.write(chunk)
                                    pbar.update(len(chunk))
                    else:
                        # 无文件大小信息，直接写入
                        for chunk in response.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                
                print(f"文件 {safe_filename} 下载成功")
                if progress_bar:
                    progress_bar.update(1)
                return True
                
            except requests.exceptions.RequestException as e:
                print(f"下载 {safe_filename} 失败 (尝试 {attempt + 1}): {e}")
                if attempt < self.max_retries - 1:
                    time.sleep(2 ** attempt)  # 指数退避
                else:
                    print(f"文件 {safe_filename} 下载失败，已达到最大重试次数")
                    if progress_bar:
                        progress_bar.update(1)
                    return False
            except Exception as e:
                print(f"下载 {safe_filename} 时发生未知错误: {e}")
                if progress_bar:
                    progress_bar.update(1)
                return False
        
        return False
    
    def download_all(self, list_only: bool = False) -> Dict[str, int]:
        """
        下载所有文件
        
        Args:
            list_only: 仅列出文件，不下载
            
        Returns:
            下载统计信息
        """
        file_list = self.get_file_list()
        
        if not file_list:
            print("未找到任何文件")
            return {"total": 0, "success": 0, "failed": 0, "skipped": 0}
        
        if list_only:
            print("\n文件列表:")
            for i, file_info in enumerate(file_list, 1):
                filename = file_info.get("filename", "unknown")
                file_id = file_info.get("id", "unknown")
                print(f"{i:3d}. {filename} (ID: {file_id})")
            return {"total": len(file_list), "success": 0, "failed": 0, "skipped": 0}
        
        # 统计信息
        stats = {"total": len(file_list), "success": 0, "failed": 0, "skipped": 0}
        
        # 创建进度条
        with tqdm(total=len(file_list), desc="下载进度") as pbar:
            for file_info in file_list:
                filename = file_info.get("filename", "unknown")
                safe_filename = self.sanitize_filename(filename)
                file_path = self.output_dir / safe_filename
                
                if file_path.exists():
                    stats["skipped"] += 1
                    pbar.update(1)
                    continue
                
                if self.download_file(file_info, pbar):
                    stats["success"] += 1
                else:
                    stats["failed"] += 1
        
        return stats
    
    def close(self):
        """关闭会话"""
        self.session.close()


def load_config(config_path: str) -> Dict:
    """
    加载配置文件
    
    Args:
        config_path: 配置文件路径
        
    Returns:
        配置字典
    """
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"加载配置文件失败: {e}")
        return {}


def save_config(config: Dict, config_path: str):
    """
    保存配置文件
    
    Args:
        config: 配置字典
        config_path: 配置文件路径
    """
    try:
        with open(config_path, 'w', encoding='utf-8') as f:
            json.dump(config, f, indent=2, ensure_ascii=False)
        print(f"配置已保存到: {config_path}")
    except Exception as e:
        print(f"保存配置文件失败: {e}")


def create_sample_config(config_path: str):
    """
    创建示例配置文件
    
    Args:
        config_path: 配置文件路径
    """
    sample_config = {
        "username": "your_username_here",
        "output_dir": "./land30m_data",
        "max_retries": 3,
        "timeout": 30
    }
    save_config(sample_config, config_path)


def parse_args() -> argparse.Namespace:
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="Land30m 全球土地利用数据下载器")
    
    # 基本参数
    parser.add_argument("--username", help="CAS Earth 用户名")
    parser.add_argument("--output-dir", default="./land30m_data", 
                       help="下载目录 (默认: ./land30m_data)")
    parser.add_argument("--config", help="配置文件路径")
    
    # 下载选项
    parser.add_argument("--list-only", action="store_true", 
                       help="仅列出文件，不下载")
    parser.add_argument("--max-retries", type=int, default=3, 
                       help="最大重试次数 (默认: 3)")
    parser.add_argument("--timeout", type=int, default=30, 
                       help="请求超时时间/秒 (默认: 30)")
    
    # 配置管理
    parser.add_argument("--create-config", help="创建示例配置文件")
    
    return parser.parse_args()


def main():
    """主函数"""
    args = parse_args()
    
    # 创建示例配置文件
    if args.create_config:
        create_sample_config(args.create_config)
        return
    
    # 加载配置
    config = {}
    if args.config:
        config = load_config(args.config)
    
    # 合并命令行参数和配置文件
    username = args.username or config.get("username")
    output_dir = args.output_dir or config.get("output_dir", "./land30m_data")
    max_retries = args.max_retries or config.get("max_retries", 3)
    timeout = args.timeout or config.get("timeout", 30)
    
    # 验证必需参数
    if not username:
        print("错误: 请提供用户名 (--username) 或配置文件 (--config)")
        print("提示: 使用 --create-config config.json 创建示例配置文件")
        return
    
    # 创建下载器
    downloader = Land30mDownloader(
        username=username,
        output_dir=output_dir,
        max_retries=max_retries,
        timeout=timeout
    )
    
    try:
        # 执行下载
        stats = downloader.download_all(list_only=args.list_only)
        
        # 显示统计信息
        if not args.list_only:
            print(f"\n下载完成!")
            print(f"总计: {stats['total']} 个文件")
            print(f"成功: {stats['success']} 个")
            print(f"失败: {stats['failed']} 个")
            print(f"跳过: {stats['skipped']} 个")
            print(f"下载目录: {output_dir}")
        
    except KeyboardInterrupt:
        print("\n用户中断下载")
    except Exception as e:
        print(f"下载过程中发生错误: {e}")
    finally:
        downloader.close()


if __name__ == "__main__":
    main()

